#include "muxer_mp4_writer.h"
#include <sys/stat.h>
#include "base_log.h"

namespace El {
namespace Muxer {

MP4Writer::MP4Writer() = default;

MP4Writer::~MP4Writer()
{
    Close();
}

void MP4Writer::SetVideoParams(int width, int height, int fps)
{
    if (IsOpened()) {
        LOG_WARN("Cannot set video params after file is opened");
        return;
    }
    videoParams_.width = width;
    videoParams_.height = height;
    videoParams_.fps = fps;
}

void MP4Writer::SetAudioParams(int sampleRate, int channels, int depth)
{
    if (IsOpened()) {
        LOG_WARN("Cannot set audio params after file is opened");
        return;
    }
    audioParams_.sampleRate = sampleRate;
    audioParams_.channels = channels;
    audioParams_.depth = depth;
}

bool MP4Writer::Open(const std::string &filePath)
{
    if (IsOpened()) {
        LOG_ERROR("MP4Writer already opened");
        return false;
    }

    filePath_ = filePath;

    // 创建输出上下文
    int ret = avformat_alloc_output_context2(&formatContext_, nullptr, nullptr, filePath.c_str());
    if (ret < 0 || !formatContext_) {
        LOG_ERROR("Failed to create output context, ret: {}", ret);
        return false;
    }

    // 打开输出文件
    ret = avio_open(&formatContext_->pb, filePath.c_str(), AVIO_FLAG_WRITE);
    if (ret < 0) {
        LOG_ERROR("Failed to open output file: {}, ret: {}", filePath, ret);
        Close();
        return false;
    }

    return true;
}

bool MP4Writer::InitializeStreams()
{
    if (isInitialized_) {
        return true;
    }

    // 初始化视频流
    if (videoParams_.width > 0 && videoParams_.height > 0) {
        videoStream_ = avformat_new_stream(formatContext_, nullptr);
        if (!videoStream_) {
            LOG_ERROR("Failed to create video stream");
            return false;
        }

        videoStream_->codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
        videoStream_->codecpar->codec_id = AV_CODEC_ID_H264; // 默认使用H264编码
        videoStream_->codecpar->width = videoParams_.width;
        videoStream_->codecpar->height = videoParams_.height;
        videoStream_->time_base = {1, 90000};
        videoStream_->r_frame_rate = {videoParams_.fps, 1};
        videoStream_->avg_frame_rate = videoStream_->r_frame_rate;
    }

    // 初始化音频流
    if (audioParams_.sampleRate > 0) {
        audioStream_ = avformat_new_stream(formatContext_, nullptr);
        if (!audioStream_) {
            LOG_ERROR("Failed to create audio stream");
            return false;
        }

        audioStream_->codecpar->codec_type = AVMEDIA_TYPE_AUDIO;
        audioStream_->codecpar->codec_id = AV_CODEC_ID_AAC;
        audioStream_->codecpar->sample_rate = audioParams_.sampleRate;
        audioStream_->codecpar->channels = audioParams_.channels;
        audioStream_->codecpar->bits_per_coded_sample = audioParams_.depth;
        audioStream_->codecpar->channel_layout = av_get_default_channel_layout(audioParams_.channels);
        audioStream_->time_base = {1, audioParams_.sampleRate};
    }

    // 写入文件头
    int ret = avformat_write_header(formatContext_, nullptr);
    if (ret < 0) {
        LOG_ERROR("Failed to write header, ret: {}", ret);
        return false;
    }

    isInitialized_ = true;
    return true;
}

bool MP4Writer::WriteFrame(const Media::MediaFramePtr &frame)
{
    if (!IsOpened()) {
        LOG_ERROR("MP4Writer not opened");
        return false;
    }

    if (!frame) {
        LOG_ERROR("Invalid frame");
        return false;
    }

    if (!InitializeStreams()) {
        return false;
    }

    // 如果有视频流，需要等待第一个I帧
    if (videoStream_ && !hasFirstIFrame_) {
        if (frame->IsVideoFrame() && frame->GetFrameType() == MEDIA_FRAME_I) {
            hasFirstIFrame_ = true;
            videoStartTime_ = frame->GetPts();
            LOG_DEBUG("First I-frame received, starting to write frames");
        } else if (frame->IsVideoFrame()) {
            return true; // 只跳过视频非I帧
        }
    } else if (!videoStream_ && !hasFirstIFrame_) {
        // 纯音频录制，不需要等待I帧
        hasFirstIFrame_ = true;
        audioStartTime_ = frame->GetPts();
        LOG_INFO("Audio only recording started");
    }

    // 根据帧类型设置对应的开始时间
    if (frame->IsVideoFrame()) {
        if (videoStartTime_ == -1) {
            videoStartTime_ = frame->GetPts();
        }
    } else {
        if (audioStartTime_ == -1) {
            audioStartTime_ = frame->GetPts();
        }
    }

    return frame->IsVideoFrame() ? WriteVideoFrame(frame) : WriteAudioFrame(frame);
}

bool MP4Writer::WriteVideoFrame(const Media::MediaFramePtr &frame)
{
    if (!videoStream_) {
        LOG_ERROR("Video stream not initialized");
        return false;
    }

    AVPacket *packet = av_packet_alloc();
    if (!packet) {
        LOG_ERROR("Failed to allocate packet");
        return false;
    }

    packet->data = frame->GetBuffer();
    packet->size = frame->GetLength();

    // 使用视频开始时间计算pts
    int64_t pts = (frame->GetPts() - videoStartTime_) * videoStream_->time_base.den / 1000000;

    if (lastVideoDts_ >= pts) {
        // 如果当前pts小于或等于上一个dts，则递增
        pts = lastVideoDts_ + 1;
    }

    packet->pts = pts;
    packet->dts = pts;
    lastVideoDts_ = pts;

    packet->stream_index = videoStream_->index;

    int ret = av_interleaved_write_frame(formatContext_, packet);
    av_packet_free(&packet);

    if (ret < 0) {
        LOG_ERROR("Failed to write video frame, ret: {}", ret);
        return false;
    }

    return true;
}

bool MP4Writer::WriteAudioFrame(const Media::MediaFramePtr &frame)
{
    if (!audioStream_) {
        LOG_ERROR("Audio stream not initialized");
        return false;
    }

    AVPacket *packet = av_packet_alloc();
    if (!packet) {
        LOG_ERROR("Failed to allocate packet");
        return false;
    }

    packet->data = frame->GetBuffer();
    packet->size = frame->GetLength();

    // 使用音频开始时间计算pts
    int64_t pts = (frame->GetPts() - audioStartTime_) * audioStream_->time_base.den / 1000000;

    if (lastAudioDts_ >= pts) {
        // 如果当前pts小于或等于上一个dts，则递增
        pts = lastAudioDts_ + 1;
    }

    packet->pts = pts;
    packet->dts = pts;
    lastAudioDts_ = pts;

    packet->stream_index = audioStream_->index;

    // LOG_INFO("Writing audio frame: pts = {}, dts = {}, size = {}", pts, packet->dts, packet->size);
    int ret = av_interleaved_write_frame(formatContext_, packet);
    av_packet_free(&packet);

    if (ret < 0) {
        char errBuf[AV_ERROR_MAX_STRING_SIZE] = {0};
        av_strerror(ret, errBuf, AV_ERROR_MAX_STRING_SIZE);
        LOG_ERROR("Failed to write audio frame, error: {}, ret: {}", errBuf, ret);
        return false;
    }

    return true;
}

void MP4Writer::Close()
{
    if (formatContext_) {
        if (isInitialized_) {
            av_write_trailer(formatContext_);
        }
        if (formatContext_->pb) {
            avio_closep(&formatContext_->pb);
        }
        avformat_free_context(formatContext_);
        formatContext_ = nullptr;
    }

    videoStream_ = nullptr;
    audioStream_ = nullptr;
    videoStartTime_ = -1;
    audioStartTime_ = -1;
    isInitialized_ = false;
    lastVideoDts_ = -1;
    lastAudioDts_ = -1;
}

uint64_t MP4Writer::GetDuration() const
{
    if (!formatContext_ || !isInitialized_) {
        return 0;
    }

    uint64_t duration = 0;

    // 使用最后一帧的 PTS 计算视频流时长
    if (videoStream_ && lastVideoDts_ >= 0) {
        duration = std::max(duration, static_cast<uint64_t>(lastVideoDts_ * av_q2d(videoStream_->time_base) * 1000000));
    }

    // 使用最后一帧的 PTS 计算音频流时长
    if (audioStream_ && lastAudioDts_ >= 0) {
        duration = std::max(duration, static_cast<uint64_t>(lastAudioDts_ * av_q2d(audioStream_->time_base) * 1000000));
    }

    return duration / 1000000;
}

/**
 * Get the current file size in bytes
 * @return File size in bytes, or 0 if file is not opened
 */
uint64_t MP4Writer::GetFileSize() const
{
    if (!formatContext_ || !formatContext_->pb) {
        return 0;
    }
    return avio_size(formatContext_->pb);
}

} // namespace Muxer
} // namespace El